from flask import Blueprint, request, jsonify, Response, stream_with_context
from services.indicator_agent_service import IndicatorAgentService
from typing import Dict, Any
import json
import asyncio

indicator_bp = Blueprint('indicator', __name__)
indicator_service = IndicatorAgentService()

@indicator_bp.route('/process-query', methods=['POST'])
def process_indicator_query():
    """
    处理指标查询的API端点
    
    请求体格式:
    {
        "query": "查询文本",
        "dark_mode": false  // 可选，默认为false
    }
    
    返回:
    {
        "parsed_result": object,
        "match_result": object,
        "dimensions": array,
        "filters": array,
        "query_params": object,
        "data": array,
        "plot_code": object,
        "figure": object,
        "status": "success" | "error",
        "message": string
    }
    """
    try:
        data = request.get_json()
        if not data or 'query' not in data:
            return jsonify({
                'status': 'error',
                'message': '缺少必要的查询参数'
            }), 400

        query = data['query']
        dark_mode = data.get('dark_mode', False)

        result = indicator_service.process_indicator_query(query, dark_mode)
        
        # 检查数据类型并进行相应的转换
        if result.get('data') is not None:
            if hasattr(result['data'], 'to_dict'):  # 如果是DataFrame
                result['data'] = result['data'].to_dict(orient='records')
            # 如果已经是list就不需要转换
            
        # 如果结果包含Plotly图形，将其转换为JSON
        if result.get('figure') is not None:
            result['figure'] = result['figure'].to_json()

        return jsonify(result)

    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': f'处理请求时发生错误: {str(e)}'
        }), 500

@indicator_bp.route('/process-query-stream', methods=['POST'])
def query_indicators_stream():
    try:
        data = request.get_json()
        query = data.get('query')
        dark_mode = data.get('dark_mode', False)

        if not query:
            return jsonify({'error': 'Query is required'}), 400

        async def generate():
            async for result in indicator_service.process_indicator_query_stream(query, dark_mode):
                # 如果result中包含figure对象,需要特殊处理,因为figure对象无法直接序列化
                result_copy = result.copy()
                if 'figure' in result_copy and result_copy['figure'] is not None:
                    result_copy['figure'] = True  # 仅标记存在figure
                else:
                    result_copy['figure'] = False
                    
                # 将DataFrame转换为JSON格式
                if 'data' in result_copy and result_copy['data'] is not None:
                    result_copy['data'] = result_copy['data'].to_json(orient='records')
                    
                yield f"data: {json.dumps(result_copy)}\n\n"

        async def run_async():
            return [chunk async for chunk in generate()]

        chunks = asyncio.run(run_async())
        
        return Response(
            (chunk for chunk in chunks),
            mimetype='text/event-stream',
            headers={
                'Cache-Control': 'no-cache',
                'Connection': 'keep-alive',
                'X-Accel-Buffering': 'no'
            }
        )

    except Exception as e:
        return jsonify({'error': str(e)}), 500